"""
python fujian1_lstm_predict.py
"""

import os
import json
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from datetime import timedelta, datetime
import tensorflow as tf
import warnings
import contextlib
import io
from PyQt5 import QtWidgets
import sys

# 还是不屏蔽吧
# # 屏蔽 TensorFlow 的不必要日志输出
# tf.get_logger().setLevel('ERROR')
# # 屏蔽所有警告
# warnings.filterwarnings('ignore')

# 文件路径
input_dir = 'fujian/fujian1/preprocessed'
output_dir = 'fujian/fujian1/predict_lstm'
loss_output_dir = 'fujian/fujian1/predict_lstm/epoch_loss'  # 新增损失输出目录

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
os.makedirs(loss_output_dir, exist_ok=True)  # 创建损失输出目录

# 预测日期范围
start_date = datetime(2023, 5, 16)
end_date = datetime(2023, 5, 30)

def create_dataset(data, time_step=1):
    X, y = [], []
    for i in range(len(data) - time_step):
        X.append(data[i:(i + time_step), 0])
        y.append(data[i + time_step, 0])
    return np.array(X), np.array(y)

class ProgressWindow(QtWidgets.QWidget):
    def __init__(self, total_files):
        super().__init__()
        self.initUI(total_files)

    def initUI(self, total_files):
        self.setWindowTitle('Processing JSON Files')
        self.setGeometry(300, 300, 400, 100)

        self.layout = QtWidgets.QVBoxLayout()

        self.progressBar = QtWidgets.QProgressBar(self)
        self.progressBar.setMaximum(total_files)
        self.layout.addWidget(self.progressBar)

        self.setLayout(self.layout)

    def update_progress(self, value):
        self.progressBar.setValue(value)

def main():
    # 隐藏所有 CPU 设备
    tf.config.set_visible_devices([], 'CPU')

    # 确保 GPU 可用
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        # 设置 GPU 记忆增长，这样程序运行时 GPU 会逐渐分配显存
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    else:
        print("No GPU available.")

    # 处理每个 JSON 文件
    json_files = [f for f in os.listdir(input_dir) if f.endswith('.json')]
    
    # 创建进度窗口
    app = QtWidgets.QApplication(sys.argv)
    progress_window = ProgressWindow(len(json_files))
    progress_window.show()

    for idx, filename in enumerate(json_files):
        with open(os.path.join(input_dir, filename), 'r') as f:
            json_data = json.load(f)

        # 提取 qty 数据
        qty_data = np.array([entry['qty'] for entry in json_data]).reshape(-1, 1)
        
        # 数据归一化
        scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = scaler.fit_transform(qty_data)

        # 创建数据集
        time_step = 10  # 可以根据需要调整
        X, y = create_dataset(scaled_data, time_step)
        X = X.reshape(X.shape[0], X.shape[1], 1)  # 转换为LSTM输入格式

        # 添加这行代码以屏蔽与 Keras 相关的 UserWarning
        warnings.filterwarnings("ignore", category=UserWarning, module="keras")

        # 构建 LSTM 模型
        model = Sequential()
        model.add(LSTM(50, return_sequences=True, input_shape=(X.shape[1], 1)))  # 修改这里
        model.add(LSTM(50))
        model.add(Dense(1))
        model.compile(optimizer='adam', loss='mean_squared_error')

        # 训练模型并记录每个 epoch 的损失
        history = model.fit(X, y, epochs=100, batch_size=32, verbose=0)

        # 将损失数据输出到文件
        loss_output_filename = os.path.join(loss_output_dir, f'loss_{filename}')
        with open(loss_output_filename, 'w') as f:
            json.dump(history.history['loss'], f)

        # 进行预测
        last_data = scaled_data[-time_step:]  # 获取最后的 time_step 个数据
        predictions = []
        for _ in range((end_date - start_date).days + 1):
            last_data = last_data.reshape((1, time_step, 1))
            
            # 使用 contextlib.redirect_stdout 重定向输出
            with contextlib.redirect_stdout(io.StringIO()):
                pred = model.predict(last_data)  # 重定向模型预测输出
                
            predictions.append(pred[0, 0])
            last_data = np.append(last_data[0][1:], pred)  # 更新 last_data

        # 反归一化
        predictions = scaler.inverse_transform(np.array(predictions).reshape(-1, 1))

        # 创建输出数据
        output_data = []
        for i in range(len(predictions)):
            date = (start_date + timedelta(days=i)).strftime('%Y-%m-%d')
            output_data.append({
                "seller_no": json_data[0]['seller_no'],
                "product_no": json_data[0]['product_no'],
                "warehouse_no": json_data[0]['warehouse_no'],
                "date": date,
                "qty": float(predictions[i][0])  # 访问 NumPy 数组的值
            })

        # 保存预测结果
        output_filename = os.path.join(output_dir, f'predicted_{filename}')
        with open(output_filename, 'w') as f:
            json.dump(output_data, f, indent=4)

        # 更新进度条
        progress_window.update_progress(idx + 1)
        QtWidgets.QApplication.processEvents()  # 更新界面

    print("预测完成，结果已保存。")
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()
